"""
Define the interface for algorithms to access relations
"""
from abc import ABC
from typing import Dict, List, Set, Tuple, Union, Any

import networkx as nx

from ..utils import dump_json
from ..utils import load_json
from ..utils import topological_sort


class Node:
    """
    The element of a graph
    """

    def __init__(self, entity: str, metric: str):
        self._entity = entity
        self._metric = metric

    @property
    def entity(self) -> str:
        """
        Entity getter
        """
        return self._entity

    @property
    def metric(self) -> str:
        """
        Metric getter
        """
        return self._metric

    def asdict(self) -> Dict[str, str]:
        """
        Serialized as a dict
        """
        return {"entity": self._entity, "metric": self._metric}

    def __eq__(self, obj: object) -> bool:
        if isinstance(obj, Node):
            return self.entity == obj.entity and self.metric == obj.metric
        return False

    def __hash__(self) -> int:
        return hash((self.entity, self.metric))

    def __repr__(self) -> str:
        return f"Node{(self.entity, self.metric)}"


class LoadingInvalidGraphException(Exception):
    """
    This exception indicates that Graph tries to load from a broken file
    """


class Graph(ABC):
    """
    The abstract interface to access relations
    """

    def __init__(self):
        self._nodes: Set[Node] = set()
        self._sorted_nodes: List[Set[Node]] = None

    def dump(self, filename: str) -> bool:
        # pylint: disable=no-self-use, unused-argument
        """
        Dump a graph into the given file

        Return whether the operation succeeds
        """
        return False

    @classmethod
    def load(cls, filename: str) -> Union["Graph", None]:
        # pylint: disable=unused-argument
        """
        Load a graph from the given file

        Returns:
        - A graph, if available
        - None, if dump/load is not supported
        - Raise LoadingInvalidGraphException if the file cannot be parsed
        """
        return None

    @property
    def nodes(self) -> Set[Node]:
        """
        Get the set of nodes in the graph
        """
        return self._nodes

    @property
    def topological_sort(self) -> List[Set[Node]]:
        """
        Sort nodes with parents first

        The graph specifies the parents of each node.
        """
        if self._sorted_nodes is None:
            self._sorted_nodes = topological_sort(
                nodes=self.nodes, predecessors=self.parents, successors=self.children
            )
        return self._sorted_nodes

    def children(self, node: Node, **kwargs) -> Set[Node]:
        """
        Get the children of the given node in the graph
        """
        raise NotImplementedError

    def children_with_edge_data(self, node: Node, **kwargs) -> Dict[Node, Dict[str, Any]]:
        """
        Get the children of the given node in the graph along with edge data
        Returns a dictionary of {child_node: edge_data}
        """
        raise NotImplementedError

    def parents(self, node: Node, **kwargs) -> Set[Node]:
        """
        Get the parents of the given node in the graph
        """
        raise NotImplementedError
        
    def parents_with_edge_data(self, node: Node, **kwargs) -> Dict[Node, Dict[str, Any]]:
        """
        Get the parents of the given node in the graph along with edge data
        Returns a dictionary of {parent_node: edge_data}
        """
        raise NotImplementedError


class MemoryGraph(Graph):
    """
    Implement Graph with data in memory
    """

    def __init__(self, graph: nx.DiGraph):
        """
        graph: The whole graph
        """
        super().__init__()
        self._graph = graph
        self._nodes.update(self._graph.nodes)

    def dump(self, filename: str) -> bool:
        nodes: List[Node] = list(self._graph.nodes)
        node_indexes = {node: index for index, node in enumerate(nodes)}
        
        # 包含边的属性数据
        edges_data = []
        for cause, effect, edge_data in self._graph.edges(data=True):
            edge_info = {
                "source": node_indexes[cause],
                "target": node_indexes[effect],
                "attributes": edge_data
            }
            edges_data.append(edge_info)
            
        data = dict(
            nodes=[node.asdict() for node in nodes], 
            edges=edges_data
        )
        dump_json(filename=filename, data=data)
        return True

    @classmethod
    def load(cls, filename: str) -> Union["MemoryGraph", None]:
        data: dict = load_json(filename=filename)
        if "nodes" not in data or "edges" not in data:
            raise LoadingInvalidGraphException(filename)
        
        nodes: List[Node] = [Node(**node) for node in data["nodes"]]
        graph = nx.DiGraph()
        graph.add_nodes_from(nodes)
        
        # 添加边并保留边的属性
        for edge_info in data["edges"]:
            source_idx = edge_info["source"]
            target_idx = edge_info["target"]
            attributes = edge_info.get("attributes", {})
            
            graph.add_edge(nodes[source_idx], nodes[target_idx], **attributes)
            
        return MemoryGraph(graph)

    def children(self, node: Node, **kwargs) -> Set[Node]:
        if not self._graph.has_node(node):
            return set()
        return set(self._graph.successors(node))
    
    def children_with_edge_data(self, node: Node, **kwargs) -> Dict[Node, Dict[str, Any]]:
        """
        获取节点的子节点以及对应边的属性数据
        返回格式: {子节点: 边属性字典}
        """
        if not self._graph.has_node(node):
            return {}
            
        result = {}
        for child in self._graph.successors(node):
            edge_data = self._graph.get_edge_data(node, child)
            result[child] = edge_data or {}
            
        return result

    def parents(self, node: Node, **kwargs) -> Set[Node]:
        if not self._graph.has_node(node):
            return set()
        return set(self._graph.predecessors(node))
    
    def parents_with_edge_data(self, node: Node, **kwargs) -> Dict[Node, Dict[str, Any]]:
        """
        获取节点的父节点以及对应边的属性数据
        返回格式: {父节点: 边属性字典}
        """
        if not self._graph.has_node(node):
            return {}
            
        result = {}
        for parent in self._graph.predecessors(node):
            edge_data = self._graph.get_edge_data(parent, node)
            result[parent] = edge_data or {}
            
        return result
